import torch
import torch.nn as nn
import torchvision
import torch.backends.cudnn as cudnn
import torch.optim
import os
import sys
import argparse
import time
import dataloader
import model
import numpy as np
from torchvision import transforms
from PIL import Image
import glob
import time



def lowlight(image_path,dev):
	if dev == "cuda":

		os.environ['CUDA_VISIBLE_DEVICES']='0'
		data_lowlight = Image.open(image_path)

		data_lowlight = (np.asarray(data_lowlight)/255.0)

		data_lowlight = torch.from_numpy(data_lowlight).float()
		data_lowlight = data_lowlight.permute(2,0,1)
		data_lowlight = data_lowlight.cuda().unsqueeze(0)

		DCE_net = model.enhance_net_nopool().cuda()
		DCE_net.load_state_dict(torch.load('snapshots/Epoch99.pth'))
	elif dev == "cpu":
		# --- 数据加载和预处理（CPU版本）---
		# 打开图像文件
		data_lowlight = Image.open(image_path)
		# 转换为 numpy 数组并归一化
		data_lowlight = np.asarray(data_lowlight) / 255.0
		# 转换为 PyTorch 张量（默认在CPU上）
		data_lowlight = torch.from_numpy(data_lowlight).float()
		# 调整维度顺序 [H, W, C] -> [C, H, W]
		data_lowlight = data_lowlight.permute(2, 0, 1)
		# 添加 batch 维度并确保在CPU上（无需显式调用.cpu()，因为默认已在CPU）
		data_lowlight = data_lowlight.unsqueeze(0)  # 等同于 .unsqueeze(0).cpu()
		# --- 模型加载（CPU版本）---
		# 实例化模型（默认在CPU上）
		DCE_net = model.enhance_net_nopool()  # 移除了 .cuda()
		# 加载权重时指定映射到CPU（关键！）
		DCE_net.load_state_dict(
			torch.load('snapshots/Epoch99.pth', map_location=torch.device('cpu'))
		)

		# --- 后续推理 ---
		# 直接使用CPU计算
		enhanced_image = DCE_net(data_lowlight)
	else:
		return

	start = time.time()
	_,enhanced_image,_ = DCE_net(data_lowlight)

	end_time = (time.time() - start)
	print(end_time)
	image_path = image_path.replace('test_data','result')
	result_path = image_path
	if not os.path.exists(image_path.replace('/'+image_path.split("/")[-1],'')):
		os.makedirs(image_path.replace('/'+image_path.split("/")[-1],''))

	torchvision.utils.save_image(enhanced_image, result_path)

if __name__ == '__main__':
# test_images
	with torch.no_grad():
		filePath = 'data/test_data/'
	
		file_list = os.listdir(filePath)
		
		for file_name in file_list:
			# test_list = glob.glob(filePath+file_name+"/*") 
			test_list = glob.glob(os.path.join(filePath+file_name, "*"))
			for image in test_list:
				# image = image
				print(image)
				lowlight(image,'cpu')

		

